import numpy as np
import torch
import torch.nn as nn


def get_hop_distance(num_node, edge, max_hop=1):
    # 创建一个形状为 (num_node, num_node) 的全零矩阵 A
    A = np.zeros((num_node, num_node))
    # 根据边的信息填充矩阵 A，使得 A[j, i] 和 A[i, j] 为 1，表示节点 i 和 j 之间有边
    for i, j in edge:
        A[j, i] = 1
        A[i, j] = 1

    # 创建一个形状为 (num_node, num_node) 的全 inf 矩阵 hop_dis，用于存储节点之间的跳数距离
    hop_dis = np.zeros((num_node, num_node)) + np.inf
    # 计算转移矩阵列表 transfer_mat，每个元素是 A 的 d 次幂
    transfer_mat = [np.linalg.matrix_power(A, d) for d in range(max_hop + 1)]
    # 检查转移矩阵中是否有非零元素，得到布尔矩阵 arrive_mat
    arrive_mat = (np.stack(transfer_mat) > 0)
    # 根据 arrive_mat 更新 hop_dis 矩阵，将到达的节点的跳数设置为 d
    for d in range(max_hop, -1, -1):
        hop_dis[arrive_mat[d]] = d
    # 返回节点之间的跳数距离矩阵
    return hop_dis


def normalize_digraph(A):
    # 计算矩阵 A 每列的和
    Dl = np.sum(A, 0)
    # 获取矩阵 A 的行数（节点数）
    num_node = A.shape[0]
    # 创建一个形状为 (num_node, num_node) 的全零矩阵 Dn
    Dn = np.zeros((num_node, num_node))
    # 对 Dn 进行对角线赋值，Dn[i, i] 为 Dl[i] 的倒数
    for i in range(num_node):
        if Dl[i] > 0:
            Dn[i, i] = Dl[i] ** (-1)
    # 计算归一化后的矩阵 AD
    AD = np.dot(A, Dn)
    # 返回归一化后的矩阵
    return AD


class Graph:
    def __init__(self, layout='openpose', strategy='uniform', max_hop=1, dilation=1, num_nodes=15):
        # 最大跳数
        self.max_hop = max_hop
        # 膨胀因子
        self.dilation = dilation
        # 节点数量
        self.num_node = num_nodes

        # 根据布局获取边的信息
        self.get_edge(layout)
        # 计算节点之间的跳数距离
        self.hop_dis = get_hop_distance(self.num_node, self.edge, max_hop=max_hop)
        # 根据策略获取邻接矩阵
        self.get_adjacency(strategy)

    def __str__(self):
        # 返回邻接矩阵 A 的字符串表示
        return str(self.A)

    def get_edge(self, layout):
        if layout == 'openpose':
            # 自连接边，即节点与自身的连接
            self_link = [(i, i) for i in range(self.num_node)]
            # 邻居连接边
            neighbor_link = [(0, 1), (1, 2), (1, 3), (2, 4), (1, 5), (3, 6), (4, 7), (5, 8), (6, 10), (5, 9), (9, 11),
                             (10, 12), (12, 13), (13, 14)]
            # 总的边信息
            self.edge = self_link + neighbor_link
            # 中心节点
            self.center = 1
        else:
            # 如果布局不存在，抛出异常
            raise ValueError("Do Not Exist This Layout.")

    def get_adjacency(self, strategy):
        # 根据膨胀因子和最大跳数，生成有效的跳数范围
        valid_hop = range(0, self.max_hop + 1, self.dilation)
        # 创建一个形状为 (num_node, num_node) 的全零矩阵 adjacency
        adjacency = np.zeros((self.num_node, self.num_node))
        # 根据有效跳数范围，填充 adjacency 矩阵
        for hop in valid_hop:
            adjacency[self.hop_dis == hop] = 1
        # 对 adjacency 矩阵进行归一化
        normalize_adjacency = normalize_digraph(adjacency)

        if strategy == 'uniform':
            # 创建一个形状为 (1, num_node, num_node) 的全零矩阵 A，并将归一化后的邻接矩阵赋值给 A[0]
            A = np.zeros((1, self.num_node, self.num_node))
            A[0] = normalize_adjacency
            self.A = A
        elif strategy == 'three_streams':
            # 自连接矩阵
            self_loop = np.eye(self.num_node)
            # 向心矩阵
            centripetal = np.zeros((self.num_node, self.num_node))
            # 离心矩阵
            centrifugal = np.zeros((self.num_node, self.num_node))

            # 根据跳数距离填充向心矩阵和离心矩阵
            for i in range(self.num_node):
                for j in range(self.num_node):
                    if self.hop_dis[i, j] == 1:
                        if i == self.center:
                            centrifugal[i, j] = 1
                        elif j == self.center:
                            centripetal[i, j] = 1

            # 对自连接矩阵、向心矩阵和离心矩阵进行归一化
            self_loop = normalize_digraph(self_loop)
            centripetal = normalize_digraph(centripetal)
            centrifugal = normalize_digraph(centrifugal)

            # 将归一化后的矩阵堆叠成一个新的矩阵 A
            A = np.stack([self_loop, centripetal, centrifugal])
            self.A = A
        elif strategy == 'adaptive':
            # 创建一个可学习的参数 A，初始值为单位矩阵的扩展形式
            self.A = nn.Parameter(torch.eye(self.num_node).unsqueeze(0), requires_grad=True)
        else:
            # 如果策略不存在，抛出异常
            raise ValueError("Do Not Exist This Strategy")